import torch
import torch.nn as nn
import numpy as np


def is_power_of_2(num:int):
    '''
    判断输入数字是否是n的幂次
    '''
    return num > 0 and num & (num - 1) == 0


class Noise:
    def __init__(self, n, len, device='cpu'):
        self.n = n
        self.len = len
        self.device = device

    def __call__(self):
        return torch.randn(self.n, self.len).to(self.device)


class Generator(nn.Module):
    def __init__(self, img_size:int=32, input_channels:int=64):
        super(Generator, self).__init__()
        assert is_power_of_2(img_size) and img_size >= 32, \
            'illegal image size: {}'.format(img_size)

        '''
        起始層：輸入是噪聲，輸出是2^m x 4 x 4的矩陣。
        中間層：每層長寬放大2倍，通道數減半。
        終止層：輸出通道數3，長寬為img_size。
        除終止層外，每一層都有Batch Normaliztion。
        終止層的激活函數是tanh，其它各層的激活函數是ReLU。
        '''
        layers = []
        layer_num = int(round(np.log2(img_size))) - 1  # 網絡的總層數
        print('生成器的網絡共有{}層'.format(layer_num))
        channels = int(2**(6 + layer_num))  # 第一層卷積的輸出通道數
        for i in range(layer_num):
            if i == 0:  # 第一層
                # 輸入是 in_channels x 1 x 1
                layers.append(nn.ConvTranspose2d(input_channels, channels, 4))
                layers.append(nn.BatchNorm2d(channels))
                layers.append(nn.ReLU(inplace=True))
                # 輸出是 channels x 4 x 4
            elif i == layer_num - 1:  # 最後一層
                # 放大兩倍
                layers.append(nn.ConvTranspose2d(channels, 3, 4, stride=2, padding=1))
                layers.append(nn.Tanh())
            else:  # 中間層
                # 放大兩倍
                layers.append(nn.ConvTranspose2d(channels, channels // 2, 4, stride=2, padding=1))
                layers.append(nn.BatchNorm2d(channels // 2))
                layers.append(nn.ReLU(inplace=True))
                # 通道數減半
                channels //= 2
        
        self.seq = nn.Sequential(*layers)


    def forward(self, x):
        assert len(x.shape) == 2 or len(x.shape) == 4
        if len(x.shape) == 2: 
            n, c = x.shape
            x = x.reshape((n, c, 1, 1))
        n, c, h, w = x.shape
        assert h == w == 1
        x = self.seq(x)
        return x.add(1).div(2)


class Discriminator(nn.Module):
    def __init__(self, img_size=32):
        super(Discriminator, self).__init__()
        assert is_power_of_2(img_size) and img_size >= 32, \
            'illegal image size: {}'.format(img_size)
        self.img_size = img_size

        '''
        起始層：輸入是圖像，其通道數是3。輸出是2^n x img_size//2 x img_size//2的矩陣。
        中間層：每層的長寬都是縮小兩倍。縮小到4x4為止。通道數加倍。激活函數是leaky ReLU。
        中間層的每一層都有Batch Normaliztion。
        最後一層：最後一層輸出維度是1x1x1。激活函數是sigmoid。       
        '''
        layers = []
        layer_num = int(round(np.log2(img_size))) - 1  # 網絡的總層數
        print('判別器的網絡共有{}層'.format(layer_num))
        channels = int(2**(3 + layer_num))  # 第一層卷積的輸出通道數
        for i in range(layer_num):
            if i == 0:  # 第一層
                # 輸入是 in_channels x img_size x img_size
                layers.append(nn.Conv2d(3, channels, 4, stride=2, padding=1))
                layers.append(nn.LeakyReLU(0.2, inplace=True))
                # 輸出是 channels x img_size/2 x img_size/2
            elif i == layer_num - 1:  # 最後一層
                # 輸入是 channels x 4 x 4
                layers.append(nn.Conv2d(channels, 1, 4))
                layers.append(nn.Sigmoid())
                # 輸出是 1 x 1 x 1
            else:  # 中間層
                # 縮小兩倍
                layers.append(nn.Conv2d(channels, channels * 2, 4, stride=2, padding=1))
                layers.append(nn.BatchNorm2d(channels * 2))
                layers.append(nn.LeakyReLU(0.2, inplace=True))
                # 通道數加倍
                channels *= 2
        
        self.seq = nn.Sequential(*layers)

    def forward(self, x):
        assert len(x.shape) == 4
        n, c, h, w = x.shape
        assert c == 3 and h == w == self.img_size
        x = self.seq(x)
        x = x.reshape(n, 1)
        return x


def simple_test():
    g = Generator()
    x = torch.rand(1, 64, 1, 1)
    assert (1, 3, 32, 32) == g(x).shape  
    d = Discriminator()
    s = d(g(x))
    assert (1, 1) == s.shape

    g = Generator(img_size=64)
    assert (1, 3, 64, 64) == g(x).shape  
    d = Discriminator(img_size=64)
    s = d(g(x))
    assert (1, 1) == s.shape


if __name__ == '__main__':
    simple_test()
